#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
火山方舟 TTS 批量语音合成脚本
参考官方 demo.py 实现，支持批量生成单词卡片语音文件

使用方法：
1. 设置环境变量：
   export VOLC_APPID="你的appid"
   export VOLC_ACCESS_TOKEN="你的access_token"  
   export VOLC_CLUSTER="你的cluster"
   export VOLC_VOICE_TYPE="你的voice_type"

2. 运行脚本：
   python3 scripts/generate_tts_volc.py

可选参数：
   --workers N        并发数量 (默认: 1)
   --overwrite        覆盖已存在文件
   --favorites        仅生成收藏单词
   --format FORMAT    音频格式 mp3/wav (默认: mp3)
   --speed SPEED      语速比例 0.5-2.0 (默认: 1.0)
   --volume VOLUME    音量比例 0.5-2.0 (默认: 1.0)  
   --pitch PITCH      音调比例 0.5-2.0 (默认: 1.0)
"""

import argparse
import base64
import json
import os
import re
import sys
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Set, Any

import requests
from tqdm import tqdm


def slugify(text: str) -> str:
    """将文本转换为安全的文件名"""
    # 移除或替换特殊字符
    text = re.sub(r'[<>:"/\\|?*]', '', text)
    text = re.sub(r'\s+', '_', text.strip())
    return text[:50]  # 限制长度


def parse_dict_json(dict_path: Path) -> List[Dict]:
    """解析 dict.json 文件，提取单词数据"""
    try:
        with open(dict_path, 'r', encoding='utf-8') as f:
            words = json.load(f)
        print(f"✓ 成功解析 {len(words)} 个单词")
        return words
        
    except FileNotFoundError:
        print(f"✗ 文件不存在: {dict_path}")
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"✗ JSON 解析失败: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"✗ 解析 dict.json 失败: {e}")
        sys.exit(1)


def load_favorites(favorites_path: Path) -> set:
    """加载收藏单词ID列表"""
    if not favorites_path.exists():
        print(f"✗ 收藏文件不存在: {favorites_path}")
        return set()
    
    try:
        favorites = json.loads(favorites_path.read_text(encoding='utf-8'))
        return set(favorites)
    except Exception as e:
        print(f"✗ 解析收藏文件失败: {e}")
        return set()


def synthesize_one(
    word_data: Dict,
    voice_config: Dict,
    output_dir: Path,
    overwrite: bool = False
) -> Tuple[bool, str]:
    """
    合成单个单词的语音
    
    Returns:
        (success: bool, message: str)
    """
    wid = word_data['id']
    text = word_data['word']
    
    # 生成文件名
    slug = slugify(f"{wid}_{text}")
    out_path = output_dir / f"{slug}.{voice_config['format']}"
    
    print(f"[准备合成] id={wid} word={text} -> {out_path.resolve()}")
    
    # 检查文件是否已存在
    if out_path.exists() and not overwrite:
        print(f"[跳过] 文件已存在: {out_path.name}")
        return True, "文件已存在"
    
    try:
        # 构建请求
        request_json = {
            "app": {
                "appid": voice_config['appid'],
                "token": "access_token",  # 固定值
                "cluster": voice_config['cluster']
            },
            "user": {
                "uid": "388808087185088"  # 固定值
            },
            "audio": {
                "voice_type": voice_config['voice_type'],
                "encoding": voice_config['format'],
                "speed_ratio": voice_config['speed'],
                "volume_ratio": voice_config['volume'],
                "pitch_ratio": voice_config['pitch'],
            },
            "request": {
                "reqid": str(uuid.uuid4()),
                "text": text,
                "text_type": "plain",
                "operation": "query",
                "with_frontend": 1,
                "frontend_type": "unitTson"
            }
        }
        
        headers = {
            "Authorization": f"Bearer;{voice_config['access_token']}",
            "Content-Type": "application/json"
        }
        
        print(f"[请求] endpoint={voice_config['api_url']} voice_type={voice_config['voice_type']} format={voice_config['format']}")
        
        # 发送请求
        response = requests.post(
            voice_config['api_url'],
            data=json.dumps(request_json),
            headers=headers,
            timeout=30
        )
        
        if response.status_code != 200:
            error_msg = f"HTTP {response.status_code}: {response.text}"
            print(f"[生成失败] id={wid} word={text} -> {out_path.resolve()} | reason=TTS 请求失败 {error_msg}")
            return False, error_msg
        
        # 解析响应
        result = response.json()
        
        if "data" not in result:
            error_msg = f"响应中无 data 字段: {result}"
            print(f"[生成失败] id={wid} word={text} -> {out_path.resolve()} | reason={error_msg}")
            return False, error_msg
        
        # 解码并保存音频数据
        audio_data = base64.b64decode(result["data"])
        
        # 确保输出目录存在
        out_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 写入文件
        with open(out_path, 'wb') as f:
            f.write(audio_data)
        
        print(f"[生成成功] id={wid} word={text} -> {out_path.resolve()} | size={len(audio_data)} bytes")
        return True, "生成成功"
        
    except Exception as e:
        error_msg = f"异常: {str(e)}"
        print(f"[生成失败] id={wid} word={text} -> {out_path.resolve()} | reason={error_msg}")
        return False, error_msg


def main():
    parser = argparse.ArgumentParser(description='火山方舟 TTS 批量语音合成')
    parser.add_argument('--workers', type=int, default=1, help='并发数量')
    parser.add_argument('--overwrite', action='store_true', help='覆盖已存在文件')
    parser.add_argument('--favorites', action='store_true', help='仅生成收藏单词')
    parser.add_argument('--format', default='mp3', choices=['mp3', 'wav'], help='音频格式')
    parser.add_argument('--speed', type=float, default=1.0, help='语速比例 (0.5-2.0)')
    parser.add_argument('--volume', type=float, default=1.0, help='音量比例 (0.5-2.0)')
    parser.add_argument('--pitch', type=float, default=1.0, help='音调比例 (0.5-2.0)')
    
    args = parser.parse_args()
    
    # 检查环境变量
    required_env = ['VOLC_APPID', 'VOLC_ACCESS_TOKEN', 'VOLC_CLUSTER', 'VOLC_VOICE_TYPE']
    missing_env = [env for env in required_env if not os.getenv(env)]
    
    if missing_env:
        print(f"✗ 缺少必需的环境变量: {', '.join(missing_env)}")
        print("\n请设置以下环境变量：")
        print("export VOLC_APPID=\"你的appid\"")
        print("export VOLC_ACCESS_TOKEN=\"你的access_token\"")
        print("export VOLC_CLUSTER=\"你的cluster\"")
        print("export VOLC_VOICE_TYPE=\"你的voice_type\"")
        sys.exit(1)
    
    # 配置参数
    voice_config = {
        'appid': os.getenv('VOLC_APPID'),
        'access_token': os.getenv('VOLC_ACCESS_TOKEN'),
        'cluster': os.getenv('VOLC_CLUSTER'),
        'voice_type': os.getenv('VOLC_VOICE_TYPE'),
        'api_url': os.getenv('VOLC_TTS_ENDPOINT', 'https://openspeech.bytedance.com/api/v1/tts'),
        'format': args.format,
        'speed': args.speed,
        'volume': args.volume,
        'pitch': args.pitch,
    }
    
    print(f"🎵 火山方舟 TTS 批量合成")
    print(f"   appid: {voice_config['appid']}")
    print(f"   cluster: {voice_config['cluster']}")
    print(f"   voice_type: {voice_config['voice_type']}")
    print(f"   format: {voice_config['format']}")
    print(f"   speed: {voice_config['speed']}")
    print(f"   volume: {voice_config['volume']}")
    print(f"   pitch: {voice_config['pitch']}")
    print(f"   workers: {args.workers}")
    print()
    
    # 路径设置
    project_root = Path(__file__).parent.parent
    dict_path = project_root / 'dict.json'
    output_dir = project_root / 'voice'
    logs_dir = project_root / 'logs'
    favorites_path = project_root / 'favorites.json'
    
    # 解析单词数据
    words = parse_dict_json(dict_path)
    
    # 过滤收藏单词
    if args.favorites:
        favorites = load_favorites(favorites_path)
        if not favorites:
            print("✗ 未找到收藏单词，退出")
            sys.exit(1)
        words = [w for w in words if w['id'] in favorites]
        print(f"✓ 过滤后剩余 {len(words)} 个收藏单词")
    
    if not words:
        print("✗ 没有需要处理的单词")
        sys.exit(1)
    
    # 创建输出目录
    output_dir.mkdir(exist_ok=True)
    logs_dir.mkdir(exist_ok=True)
    
    # 批量合成
    success_count = 0
    skip_count = 0
    fail_count = 0
    failures = []
    
    print(f"🚀 开始批量合成 {len(words)} 个单词...")
    
    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        # 提交任务
        future_to_word = {
            executor.submit(synthesize_one, word, voice_config, output_dir, args.overwrite): word
            for word in words
        }
        
        # 处理结果
        with tqdm(total=len(words), desc="合成进度") as pbar:
            for future in as_completed(future_to_word):
                word = future_to_word[future]
                try:
                    success, message = future.result()
                    if success:
                        if "已存在" in message:
                            skip_count += 1
                        else:
                            success_count += 1
                    else:
                        fail_count += 1
                        failures.append({
                            'id': word['id'],
                            'word': word['word'],
                            'error': message
                        })
                except Exception as e:
                    fail_count += 1
                    failures.append({
                        'id': word['id'],
                        'word': word['word'],
                        'error': str(e)
                    })
                
                pbar.update(1)
    
    # 保存失败记录
    if failures:
        failure_log = logs_dir / 'tts_failures.log'
        with open(failure_log, 'w', encoding='utf-8') as f:
            json.dump(failures, f, ensure_ascii=False, indent=2)
        print(f"✗ 失败记录已保存到: {failure_log}")
    
    # 输出统计
    print(f"\n📊 合成完成:")
    print(f"   ✅ 成功: {success_count}")
    print(f"   ⏭️  跳过: {skip_count}")
    print(f"   ❌ 失败: {fail_count}")
    print(f"   📁 输出目录: {output_dir.resolve()}")
    
    if fail_count > 0:
        print(f"\n💡 如需重试失败项，请检查环境变量配置或网络连接")
        sys.exit(1)
    else:
        print(f"\n🎉 所有语音文件生成完成！")


if __name__ == '__main__':
    main()